﻿using System;
using System.Collections.Generic;
//using System.Linq;
//using System.Text;
using Allegro.Framework;
using Allegro.Integrator;
//using Cosmologica.Framework;

namespace Allegro.Integrator
{
    public class BulirschStoerIntegrator : IODEIntegrator
    {
        const int KMAXX = 8; // Maximum row number used in the extrapolation
        const int IMAXX = KMAXX + 1;
        const double SAFE1 = 0.25; // Safety factors
        const double SAFE2 = 0.7;
        const double REDMAX = 1.0e-5; // Max. factor for stepsize reduction
        const double REDMIN = 0.7; // Min. factor for stepsize reduction
        const double TINY = 1.0e-30; // Prevents division by zero
        const double SCALMX = 0.1; // 1/SCALMX is the max. factor by which a stepsize can be increased

        int _n;
        private double _accuracy;
        bool[] _integrated; // Flag for each parameter. If true, variable i was integrated.
        bool[] _computed; // Flag for each parameter. If true, variable i was computed directly
        bool _first = true;
        int _kmax;
        int _kopt;
        double _epsold = -1.0;
        double _xnew;
        double[] _a;
        double[,] _alf;
        int[] _nseq;
        double[,] _d;
        double[] _err;
        double[] _x;
        double[] _ysav;
        double[] _yerr;
        double[] _yseq;
        double[] _dydx;
        double[] _yscal;
        double[] _c;

        // Midpoint method data
        double[] _ym;
        double[] _yn;

        public BulirschStoerIntegrator(int dimension)
        {
            if (dimension <= 0)
                throw new ApplicationException("Dim. must be > 0");

            _n = dimension;
            _a = new double[IMAXX + 1];
            _alf = new double[KMAXX + 1, KMAXX + 1];
            _nseq = new int[] { 0, 2, 4, 6, 8, 10, 12, 14, 16, 18 };
            _d = new double[_n, KMAXX + 1];
            _err = new double[KMAXX + 1];
            _x = new double[KMAXX + 1];
            _yerr = new double[_n];
            _ysav = new double[_n];
            _yseq = new double[_n];
            _dydx = new double[_n];
            _yscal = new double[_n];
            _c = new double[_n];

            _ym = new double[_n];
            _yn = new double[_n];
            _integrated = new bool[_n];
            _computed = new bool[_n];
        }


        #region IODEIntegrator Members

        public double Accuracy
        {
            get
            {
                return _accuracy;
            }
            set
            {
                _accuracy = value;
            }
        }

        /// <summary>
        /// Bulirsch-Stoer step with monitoring of local truncation error to ensure accuracy and adjust stepsize. Input are the dependent
        /// variable vector y[1..nv] and its derivative dydx[1...nv] at the starting value of the independent variable x. Also input are the stepsize
        /// to be attempted, htry, the required accuracy eps, and the vector yscal[1...nv] against which the error is scaled. On output, y and x are 
        /// replaced by their new values. hdid is the stepsize that was actually accomplished, and hnext is the estimated next stepsize. derivs is
        /// the user-supplied routine that computes the right-hand side derivatives. Be sure to set htry on successive steps to the value of hnext
        /// returned from the previous step, as is the case if the routine is called by BasicODESolver.Solve.
        /// </summary>
        /// <param name="x"></param>
        /// <param name="y"></param>
        /// <param name="htry"></param>
        /// <param name="hdid"></param>
        /// <param name="hnext"></param>
        /// <param name="f"></param>
        /// <param name="V"></param>
        public void Integrate(ref double xx, ref double[] y, double htry, out double hdid, out double hnext, Derivative derivs, ComputeValues values)
        {
            int km = 0, kBreak = 0;
            double eps1, errmax, fact, h, red = 0.0, scale = 1.0, work, wrkmin, xest;
            bool reduct;
            bool exitflag = false;
            double eps = Accuracy;
            errmax = 1.0;
            hnext = hdid = 0.0;

            bool isComplete = derivs(xx, y, ref _dydx, ref _integrated);
            if (!isComplete && values != null)
                isComplete = values(xx, y, ref _yseq, _integrated, ref _computed);

            h = htry; // Set stepsize to initial trial value
            for (int i = 0; i < _n; i++)
            {
                // Scaling used to monitor accuracy. This general-purpose choice can be modified if need be.
                _yscal[i] = Math.Abs(y[i]) + Math.Abs(_dydx[i] * h) + TINY;
            }

            if (eps != _epsold)
            {
                hnext = _xnew = -1.0e29; // A new tolerance, so reinitialize "impossible" values
                eps1 = SAFE1 * eps;
                _a[1] = _nseq[1] + 1;
                for (int k = 1; k <= KMAXX; k++)
                {
                    _a[k + 1] = _a[k] + _nseq[k + 1]; // Compute work coefficients Ak
                }
                for (int iq = 2; iq <= KMAXX; iq++)
                {
                    for (int k = 1; k < iq; k++)
                    {
                        _alf[k, iq] = Math.Pow(eps1, (_a[k + 1] - _a[iq + 1]) / ((_a[iq + 1] - _a[1] + 1.0) * (2 * k + 1)));
                    }
                }
                _epsold = eps;
                // Determine optimal row number for convergence
                for (int kopt = 2; kopt < KMAXX; kopt++)
                {
                    _kopt = kopt;
                    if (_a[kopt + 1] > _a[kopt] * _alf[kopt - 1, kopt])
                        break;
                    _kmax = kopt;
                }
            }
            h = htry;
            for (int i = 0; i < _n; i++)
            {
                _ysav[i] = y[i]; // Save the starting values
            }
            if (xx != _xnew || h != hnext)
            {
                // A new stepsize or a new integration: re-establish the order window
                _first = true;
                _kopt = _kmax;
            }
            reduct = false;
            for (; ; )
            {
                for (int k = 1; k <= _kmax; k++)
                {
                    kBreak = k;
                    // Evaluate the sequence of modified midpoint integrations
                    _xnew = xx + h;
                    if (_xnew == xx)
                        ErrorHandler.Error("Stepsize underflow in BulirschStoerIntegrator.Integrate");
                    // Mid-point method
                    _ModifiedMidPointMethod(_ysav, _dydx, _n, xx, h, _nseq[k], ref _yseq, derivs, values);
                    xest = Math.Sqrt(h / _nseq[k]); // Squared, since error series is even
                    // Perform extrapolation
                    _PolynomialExtrapolation(k, xest, _yseq, ref y, ref _yerr);
                    if (k != 1)
                    {
                        // Compute normalized error estimate err(k)
                        errmax = TINY;
                        for (int i = 0; i < _n; i++)
                        {
                            errmax = Math.Max(errmax, Math.Abs(_yerr[i] / _yscal[i]));
                        }
                        errmax /= eps; // Scale error relative to tolerance
                        km = k - 1;
                        _err[km] = Math.Pow(errmax / SAFE1, 1.0 / ((double)(2 * km + 1)));
                    }
                    if (k != 1 && (k >= _kopt - 1 || _first))
                    {
                        // In order window
                        if (errmax < 1.0)
                        {
                            // Converged
                            exitflag = true;
                            break;
                        }
                        if (k == _kmax || k == _kopt + 1)
                        {
                            // Check for possible stepsize reduction
                            red = SAFE2 / _err[km];
                            break;
                        }
                        else if (k == _kopt && _alf[_kopt - 1, _kopt] < _err[km])
                        {
                            red = 1.0 / _err[km];
                            break;
                        }
                        else if (_kopt == _kmax && _alf[km, _kmax - 1] < _err[km])
                        {
                            red = _alf[km, _kmax - 1] * SAFE2 / _err[km];
                            break;
                        }
                        else if (_alf[km, _kopt] < _err[km])
                        {
                            red = _alf[km, _kopt - 1] / _err[km];
                            break;
                        }
                    }
                }
                if (exitflag)
                    break;
                red = Math.Min(red, REDMIN); // Reduce stepsize by at least REDMIN and at most REDMAX
                red = Math.Max(red, REDMAX);
                h *= red;
                reduct = true;
            }
            xx = _xnew;
            hdid = h;
            _first = false;
            wrkmin = 1.0e35;
            // Compute optimal row for convergence and corresponding stepsize
            for (int kk = 1; kk <= km; kk++)
            {
                fact = Math.Max(_err[kk], SCALMX);
                work = fact * _a[kk + 1];
                if (work < wrkmin)
                {
                    scale = fact;
                    wrkmin = work;
                    _kopt = kk + 1;
                }
            }
            hnext = h / scale;
            if (_kopt >= kBreak && _kopt != _kmax && !reduct)
            {
                // Check for possible order increase, but not if stepsize was just reduced
                fact = Math.Max(scale / _alf[_kopt - 1, _kopt], SCALMX);
                if (_a[_kopt + 1] * fact <= wrkmin)
                {
                    hnext = h / fact;
                    _kopt++;
                }
            }

        }


        #endregion

        // Private methods
        void _PolynomialExtrapolation(int iest, double xest, double[] yest, ref double[] yz, ref double[] dy)
        {
            double q, f2, f1, delta;

            _x[iest] = xest; // Save current independent variable
            for (int j = 0; j < _n; j++)
            {
                dy[j] = yz[j] = yest[j];
            }
            if (iest == 1)
            {
                for (int j = 0; j < _n; j++)
                {
                    _d[j, 1] = yest[j];
                }
            }
            else
            {
                for (int j = 0; j < _n; j++)
                {
                    _c[j] = yest[j];
                }
                for (int k1 = 1; k1 < iest; k1++)
                {
                    delta = 1.0 / (_x[iest - k1] - xest);
                    f1 = xest * delta;
                    f2 = _x[iest - k1] * delta;
                    for (int j = 0; j < _n; j++)
                    {
                        // Propagate tableau 1 diagonal more
                        q = _d[j, k1];
                        _d[j, k1] = dy[j];
                        delta = _c[j] - q;
                        dy[j] = f1 * delta;
                        _c[j] = f2 * delta;
                        yz[j] += dy[j];
                    }
                }
                for (int j = 0; j < _n; j++)
                {
                    _d[j, iest] = dy[j];
                }
            }
        }

        /// <summary>
        /// Modified midpoint step. At xs, input the dependent variable vector y[1...nvar] and its derivative vector dydx[1...nvar]. 
        /// Also input is htot, the total step to be made, and nstep, the number of substeps to be used. The output is returned as
        /// yout[1..nvar], which need not be a distinct array from y; if it is distinct, however, then y and dydx are returned
        /// unmanaged.
        /// </summary>
        /// <param name="y"></param>
        /// <param name="dydx"></param>
        /// <param name="xs"></param>
        /// <param name="htot"></param>
        /// <param name="nstep"></param>
        /// <param name="yout"></param>
        /// <param name="derivs"></param>
        void _ModifiedMidPointMethod(double[] y, double[] dydx, int nvar, double xs, double htot, int nstep, ref double[] yout, Derivative derivs, ComputeValues values)
        {
            double x, swap, h2, h;

            h = htot / nstep; // Stepsize this trip
            for (int i = 0; i < _n; i++)
            {
                _ym[i] = y[i];
                _yn[i] = y[i] + h * dydx[i];
            }
            x = xs + h;
            // Will use yout for temporary storage of derivatives
            bool isComplete = derivs(x, _yn, ref yout, ref _integrated);
            if (!isComplete && values != null)
                isComplete = values(x, _yn, ref yout, _integrated, ref _computed);

            h2 = 2.0 * h;
            for (int n = 2; n <= nstep; n++)
            {
                for (int i = 0; i < nvar; i++)
                {
                    swap = _ym[i] + h2 * yout[i];
                    _ym[i] = _yn[i];
                    _yn[i] = swap;
                }
                x += h;
                isComplete = derivs(x, _yn, ref yout, ref _integrated);
                if (!isComplete && values != null)
                    isComplete = values(x, _yn, ref yout, _integrated, ref _computed);
            }
            // Last step
            for (int i = 0; i < _n; i++)
            {
                if (_integrated[i])
                {
                    yout[i] = 0.5 * (_ym[i] + _yn[i] + h * yout[i]);
                }
            }
        }
    }
}
